# Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
# generator_ascend_generate_grouped_matmul_add.py
import random
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig


@GENERATOR_REGISTRY.register("ascend_generate_grouped_matmul_add")
class GroupedMatmulAddGenerator(CaseGenerator):

    def __init__(self, config):
        super().__init__(config)
        self.tensor_dim = 0
        self.range_is_null = False


    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        #矩阵乘条件
        x_shape = case_config.inputs[0].shape
        w_shape = case_config.inputs[1].shape
        group_list_shape = case_config.inputs[2].shape
        y_shape = case_config.inputs[3].shape
        x_shape[0] = random.randint(1,2048)
        w_shape[0] = x_shape[0]
        # case_config.inputs[1].shape = w_shape
        group_list_shape[0] = random.randint(1, min(x_shape[0], 128))

        y_shape[0] = x_shape[1] * group_list_shape[0]
        y_shape[1] = w_shape[1]
        #左右矩阵类型相同
        x_dtype = case_config.inputs[0].dtype
        case_config.inputs[1].dtype = x_dtype
        case_config.inputs[2].dtype = x_dtype

        
        #groupList
        case_config.inputs[2].range_values = [1, x_shape[0]]
        # case_config.inputs[3].range_values = [random.randint(1,x_shape[0])]

        return case_config
    